from PIL import Image
import torch
from torchvision import transforms

IMG_SIZE = 640
image_file = '/Users/summy/project/python/parttime/ocr/car.jpeg'

# PIL图像转tensor
pil_to_tensor = transforms.Compose([
	transforms.Resize((IMG_SIZE, IMG_SIZE)),  # PIL图像尺寸统一
	transforms.ToTensor()  # PIL图像转tensor, (H,W,C) ->（C,H,W）,像素值[0,1]
])

# tensor转PIL图像
tensor_to_pil = transforms.Compose([
	transforms.Lambda(lambda t: t * 255),  # 像素还原
	transforms.Lambda(lambda t: t.type(torch.uint8)),  # 像素值取整
	transforms.ToPILImage(),  # tensor转回PIL图像, (C,H,W) -> (H,W,C)
])

affine_transform = transforms.Compose([
	transforms.Resize((IMG_SIZE, IMG_SIZE)),  # PIL图像尺寸统一
	transforms.RandomAffine(degrees=(-10, 10), translate=(0.1, 0.1), scale=(0.9, 1.1), shear=(-10, 10)),
	transforms.ToTensor()  # PIL图像转tensor, (H,W,C) ->（C,H,W）,像素值[0,1]
])

if __name__ == '__main__':
	img = Image.open(image_file)
	print('原图:', img.size)
	img_tensor = pil_to_tensor(img)
	print('tensor:', img_tensor.size())
	tensor_to_pil(img_tensor).show()
	
	img_tensor = affine_transform(img)
	tensor_to_pil(img_tensor).show()
